package org.example.recommender

import org.example.scala.model.Movie
import org.example.scala.model.MovieRating
import org.example.scala.model.Tag
import org.example.scala.model.MongoConfig
import org.example.scala.model.ESConfig
import org.example.java.model.Constant._

import com.mongodb.casbah.commons.MongoDBObject
import com.mongodb.casbah.{MongoClient, MongoClientURI}
import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest
import org.elasticsearch.action.admin.indices.exists.indices.IndicesExistsRequest
import org.elasticsearch.common.settings.Settings
import org.elasticsearch.common.transport.InetSocketTransportAddress
import org.elasticsearch.transport.client.PreBuiltTransportClient

import java.net.InetAddress


object DataLoader {

  //    定义常量
  //  由于是硬编码，记得修改
//  val MOVIE_DATA_PATH = "D:\\MovieRecommendSystem\\recommender\\DataLoader\\src\\main\\resources\\movies.csv"
//  val RATING_DATA_PATH = "D:\\MovieRecommendSystem\\recommender\\DataLoader\\src\\main\\resources\\ratings.csv"
//  val TAG_DATA_PATH = "D:\\MovieRecommendSystem\\recommender\\DataLoader\\src\\main\\resources\\tags.csv"

  //  val MONGODB_MOVIE_COLLECTION = "Movie"
  //  val MONGODB_RATING_COLLECTION = "Rating"
  //  val MONGODB_TAG_COLLECTION = "Tag"

//  val ES_MOVIE_INDEX = "Movie"

  def main(args: Array[String]): Unit = {

    if(args.length != 7){
      System.err.println("Usage: java -jar dataloader.jar <mongo_server> <es_http_server> <es_transport_server> <es_cluster_name> <movie_data_path> <rating_data_path> <tag_data_path>\n"
      + "<mongo_server> is the mongo server to connect, eg. localhost:27018\n"
      + "<es_http_server>is elastic search http servers eg. localhost:9200,localhost1:9200\n"
      + "<es_transport_server>is elastic search transport servers eg. localhost:9300,localhost1:9300\n"
      + "<es_cluster_name>\n"
      + "<movie_data_path>\n"
      + "<rating_data_path>\n"
      + "<tag_data_path>\n"
      )
      System.exit(1)
    }

    val mongo_server = args(0)
    val es_http_server = args(1)
    val es_transport_server = args(2)
    val es_cluster_name = args(3)

    val movie_data_path = args(4)
    val rating_data_path = args(5)
    val tag_data_path = args(6)

    //    用到的配置参数
    val config = Map(
      "spark.cores" -> "local[*]",
      //      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      //      为了不和windows本地的有冲突 改了端口转发
      "mongo.uri" -> ("mongodb://localhost:27018/" + MONGO_DATABASE),
      "mongo.db" -> "recommender",
      "es.httpHosts" -> "localhost:9200",
      "es.transportHosts" -> "localhost:9300",
      "es.index" -> ES_INDEX,
      //      "es.cluster.name" -> "elasticsearch"
      //      改回http://localhost:9200/的   "cluster_name"
      "es.cluster.name" -> "es-cluster"
    )

    //    创建一个sparkConf
    val sparkConf = new SparkConf().setMaster(config("spark.cores")).setAppName("DataLoader")

    //    创建sparkSession
    val spark = SparkSession.builder().config(sparkConf).getOrCreate()

    import spark.implicits._

    //    加载数据 RDD => DataFrame
    val movieRDD = spark.sparkContext.textFile(movie_data_path)
    val movieDF = movieRDD.map(
      item => {
        val attr = item.split("\\^")
        Movie(attr(0).toInt, attr(1).trim, attr(2).trim, attr(3).trim, attr(4).trim, attr(5).trim, attr(6).trim, attr(7).trim, attr(8).trim, attr(9).trim)
      }
    ).toDF()
    val ratingRDD = spark.sparkContext.textFile(rating_data_path)
    val ratingDF = ratingRDD.map(item => {
      val attr = item.split(",")
      MovieRating(attr(0).toInt, attr(1).toInt, attr(2).toDouble, attr(3).toInt)
    }).toDF()
    val tagRDD = spark.sparkContext.textFile(tag_data_path)
    val tagDF = tagRDD.map(item => {
      val attr = item.split(",")
      Tag(attr(0).toInt, attr(1).toInt, attr(2).trim, attr(3).toInt)
    }).toDF()


    //    将数据保存到MongoDB
    // 声明了一个 mongo 配置的隐式参数
    implicit val mongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))
    storeDataInMongoDB(movieDF, ratingDF, tagDF)

    //    数据预处理
    // newTags：一个电影多个tag 处理成 mid,tag1|tag2|tag3....
    // movieWithTagsDF：Movie数据集 + newTags
    import org.apache.spark.sql.functions._
    val newTag = tagDF.groupBy($"mid")
      .agg(concat_ws("|", collect_set($"tag"))
        .as("tags"))
      .select("mid", "tags")
    // 需要将处理后的 Tag 数据，和 Movie 数据融合，产生新的 Movie 数据
    val movieWithTagsDF = movieDF.join(newTag, Seq("mid", "mid"), "left")

    //    将数据DF保存到ES
    // 声明了一个 ES 配置的隐式参数
    implicit val esConfig = ESConfig(config.get("es.httpHosts").get, config.get("es.transportHosts").get, config.get("es.index").get, config.get("es.cluster.name").get)
    storeDataInES(movieWithTagsDF)

    spark.stop()
  }

  def storeDataInMongoDB(movieDF: DataFrame, ratingDF: DataFrame, tagDF: DataFrame)(implicit mongoConfig: MongoConfig): Unit = {
    //    新建一个mongodb的连接
    val mongoClient = MongoClient(MongoClientURI(mongoConfig.uri))

    //    如果mongodb中已经有相应的数据库，先删除
    mongoClient(mongoConfig.db)(MONGO_MOVIE_COLLECTION).dropCollection()
    mongoClient(mongoConfig.db)(MONGO_RATING_COLLECTION).dropCollection()
    mongoClient(mongoConfig.db)(MONGO_TAG_COLLECTION).dropCollection()

    //    将DF数据写入对应的mongo表中
    movieDF.write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGO_MOVIE_COLLECTION)
      .mode("overwrite")
      .format(MONGO_DRIVER_CLASS)
      .save()
    ratingDF.write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGO_RATING_COLLECTION)
      .mode("overwrite")
      .format(MONGO_DRIVER_CLASS)
      .save()
    tagDF.write
      .option("uri", mongoConfig.uri)
      .option("collection", MONGO_TAG_COLLECTION)
      .mode("overwrite")
      .format(MONGO_DRIVER_CLASS)
      .save()

    //对数据表建索引
    mongoClient(mongoConfig.db)(MONGO_MOVIE_COLLECTION).createIndex(MongoDBObject("mid" -> 1))
    mongoClient(mongoConfig.db)(MONGO_RATING_COLLECTION).createIndex(MongoDBObject("uid" -> 1))
    mongoClient(mongoConfig.db)(MONGO_RATING_COLLECTION).createIndex(MongoDBObject("mid" -> 1))
    mongoClient(mongoConfig.db)(MONGO_TAG_COLLECTION).createIndex(MongoDBObject("uid" -> 1))
    mongoClient(mongoConfig.db)(MONGO_TAG_COLLECTION).createIndex(MongoDBObject("mid" -> 1))

    //关闭 MongoDB 的连接
    mongoClient.close()
  }

  def storeDataInES(movieDF: DataFrame)(implicit eSConfig: ESConfig): Unit = {
    //新建一个配置
    val settings: Settings = Settings.builder().put("cluster.name", eSConfig.clustername).build()
    //新建一个 ES 的客户端
    val esClient = new PreBuiltTransportClient(settings)

    //需要将 TransportHosts 添加到 esClient 中
    val REGEX_HOST_PORT = "(.+):(\\d+)".r
    eSConfig.transportHosts.split(",").foreach {
      case REGEX_HOST_PORT(host: String, port: String) => {
        esClient.addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName(host), port.toInt))
      }
    }
    //需要清除掉 ES 中遗留的数据
    if (esClient.admin().indices().exists(new IndicesExistsRequest(eSConfig.index)).actionGet().isExists) {
      esClient.admin().indices().delete(new DeleteIndexRequest(eSConfig.index))
    }
    esClient.admin().indices().create(new CreateIndexRequest(eSConfig.index))
    //将数据写入到 ES 中
    movieDF
      .write
      .option("es.nodes", eSConfig.httpHosts)
      .option("es.http.timeout", "100m")
      .option("es.mapping.id", "mid")
      .mode("overwrite")
      .format(ES_DRIVER_CLASS)
      .save(eSConfig.index + "/" + ES_TYPE)
  }
}
